#include "AVLTree.h"
#include <stdio.h>
#include <stdlib.h>
#include <assert.h>
#include <time.h>
#include <iostream>
#define		max(a, b)	(((a) > (b)) ? (a) : (b))
#define		min(a, b)	(((a) > (b)) ? (b) : (a))

AVLTree::AVLTree(void) : mRoot(NULL)
{
}

AVLTree::~AVLTree(void)
{
	pTreeNode temp = mRoot;
	while (temp)
	{
		if (temp->left)
			temp = temp->left;
		else if (temp->right)
			temp = temp->right;
		else
		{
			mRoot = temp->parent;
			if (mRoot)
			{
				(mRoot->left == temp ? mRoot->left : mRoot->right) = NULL;
			}
			delete temp;
			temp = mRoot;
		}
	}
}

void AVLTree::insert(pNodeData data) 
{
	if (!data)
		return;

	pTreeNode pNode = new TreeNode();
	pNode->parent = NULL;
	pNode->left = NULL;
	pNode->right = NULL;
	pNode->h = 1;
	pNode->bf = 0;
	pNode->m_pData = data;
	if (mRoot)
	{
		if (bin_insert(mRoot, pNode))
			adjust_tree(pNode->parent);
		else
			delete pNode;
	}
	else
	{
		mRoot = pNode;
	}
}

bool AVLTree::bin_insert(pTreeNode root, pTreeNode pNode)
{
	int r = root->m_pData->compare(pNode->m_pData);
	if (r == 1)
	{
		if (root->left)
			return bin_insert(root->left, pNode);
		else
		{
			root->left = pNode;
			pNode->parent = root;
		}
	}
	else if (r == -1)
	{
		if (root->right)
			return bin_insert(root->right, pNode);
		else
		{
			root->right = pNode;
			pNode->parent = root;
		}
	}
	else
	{
		return false;
	}
	return true;
}

void AVLTree::adjust_tree(pTreeNode sub_root)
{
	while (sub_root)
	{
		int oldHeight = sub_root->h;
		calHeight(sub_root);
		calBalanceFactor(sub_root);
		if (sub_root->bf > 1)
		{
			if (sub_root->left->bf >= 0)
			{
				sub_root = sub_root->left;
				right_rotate(sub_root);
			}
			else
			{
				sub_root = sub_root->left->right;
				left_right_rotate(sub_root);
			}
		}
		else if (sub_root->bf < -1)
		{
			if (sub_root->right->bf > 0)
			{
				sub_root = sub_root->right->left;
				right_left_rotate(sub_root);
			}
			else
			{
				sub_root = sub_root->right;
				left_rotate(sub_root);
			}
		}
		if (sub_root->h == oldHeight)
			break;

		sub_root = sub_root->parent;
	}
}

void AVLTree::calBalanceFactor(pTreeNode node)
{
	node->bf = (node->left ? node->left->h : 0) - (node->right ? node->right->h : 0);
}

void AVLTree::calHeight(pTreeNode node)
{
	node->h = max(node->left ? node->left->h : 0, node->right ? node->right->h : 0) + 1;
}

void AVLTree::deleteData(pNodeData pData)
{
	pTreeNode node = search(pData);
	if (!node)
	{
		return;
	}
	
	pTreeNode adjust_root = NULL;
	if (!node->left && !node->right)
	{
		if (node->parent)
		{
			adjust_root = node->parent;
			(node->parent->left == node ? node->parent->left : node->parent->right) = NULL;
		}
		else
		{
			mRoot = NULL;
		}
	}
	else if (node->left && node->right)
	{
		pTreeNode del_node;
		if (node->bf > 0)
		{
			for (del_node = node->left; del_node->right; del_node = del_node->right);
			if (del_node->parent == node)
			{
				node->left = del_node->left;
				if (del_node->left)
					del_node->left->parent = node;
				adjust_root = node;
			}
			else
			{
				del_node->parent->right = del_node->left;
				if (del_node->left)
					del_node->left->parent = del_node->parent;
				adjust_root = del_node->parent;
			}
		}
		else
		{
			for (del_node = node->right; del_node->left; del_node = del_node->left);
			if (del_node->parent == node)
			{
				node->right = del_node->right;
				if (del_node->right)
					del_node->right->parent = node;
				adjust_root = node;
			}
			else
			{
				del_node->parent->left = del_node->right;
				if (del_node->right)
					del_node->right->parent = del_node->parent;                                                                                                                                                                                
				adjust_root = del_node->parent;
			}
		}
		node->m_pData = del_node->m_pData;
		node = del_node;
	}
	else
	{
		pTreeNode child = node->left ? node->left : node->right;
		child->parent = node->parent;
		if (node->parent)
		{
			(node->parent->left == node ? node->parent->left : node->parent->right) = child;
			adjust_root = node->parent;
		}
		else
		{
			mRoot = child;
		}
	}

	delete node;
	adjust_tree(adjust_root);
}

pTreeNode AVLTree::search(pNodeData pData)
{
	pTreeNode node = mRoot;
	while (node)
	{
		int r = node->m_pData->compare(pData);
		if (r == -1)
		{
			node = node->right;
		}
		else if (r == 1)
		{
			node = node->left;
		}
		else if (r == 0)
		{
			break;
		}
		else
		{
			return NULL;
		}
	}
	return node;
}

void AVLTree::left_rotate(pTreeNode pivot)
{
	pTreeNode root = pivot->parent;
	assert(root && root->right == pivot);
	
	pivot->parent = root->parent;
	if (root->parent)
	{
		if (root->parent->left == root)
		{
			root->parent->left = pivot;	
		}
		else
		{
			root->parent->right = pivot;
		}
	}
	else
	{
		mRoot = pivot;
	}
	
	root->parent = pivot;
	root->right = pivot->left;
	if (pivot->left)
		pivot->left->parent = root;
	pivot->left = root;

	calHeight(root);
	calBalanceFactor(root);
	calHeight(pivot);
	calBalanceFactor(pivot);
}

void AVLTree::right_rotate(pTreeNode pivot)
{
	pTreeNode root = pivot->parent;
	assert(root && root->left == pivot);

	pivot->parent = root->parent;
	if (root->parent)
	{
		if (root->parent->left == root)
		{
			root->parent->left = pivot;	
		}
		else
		{
			root->parent->right = pivot;
		}
	}
	else
	{
		mRoot = pivot;
	}

	root->parent = pivot;
	root->left = pivot->right;
	if (pivot->right)
		pivot->right->parent = root;
	pivot->right = root;

	calHeight(root);
	calBalanceFactor(root);
	calHeight(pivot);
	calBalanceFactor(pivot);
}

void AVLTree::left_right_rotate(pTreeNode pivot)
{
	left_rotate(pivot);
	right_rotate(pivot);
}

void AVLTree::right_left_rotate(pTreeNode pivot)
{
	right_rotate(pivot);
	left_rotate(pivot);
}

bool _checkAVLTree(pTreeNode root)
{
	if (!root)
		return true;
	
	int lh = 0;
	if (_checkAVLTree(root->left))
	{
		lh = root->left ? root->left->h : 0;
	}
	else
	{
		return false;
	}

	int rh = 0;
	if (_checkAVLTree(root->right))
	{
		rh = root->right ? root->right->h : 0;
	}
	else
	{
		return false;
	}
	
	bool res = root->h == max(lh, rh) + 1 && root->bf == lh - rh && root->bf >= -1 && root->bf <= 1;
	return res;
}

bool checkAVLTree(AVLTree& tree)
{
	return _checkAVLTree(tree.mRoot);
}

int unit_test()
{
	const int size = 20000;
	NodeData datas[size];
	for (int i = 0; i < size; ++i)
	{
		datas[i].mValue = i;
	}

	srand((unsigned)time(NULL));
	for (int i = 0; i < size; ++i)
	{
		int index = i + ((unsigned) rand()) % (size - i);
		NodeData temp = datas[index];
		datas[index] = datas[i];
		datas[i] = temp;
	}
	
	std::cout << "begin construct avl tree." << std::endl;
	AVLTree tree;
	for (int i = 0; i < size; ++i)
	{
		tree.insert(datas + i);
		if (!checkAVLTree(tree))
		{
			std::cout << "construct avl-tree failed" << std::endl;
			return 0;
		}
	}
	
	std::cout << "begin searching test." << std::endl;
	for (int i = 0; i < size; ++i)
	{
		if (!tree.search(datas + i))
		{
			std::cout << "search error 1 - " << i << std::endl;
		}
	}

	NodeData data0;
	data0.mValue = size;
	if (tree.search(&data0))
	{
		std::cout << "search error 2" << std::endl;
		return 0;
	}
	
	std::cout << "begin deleting test." << std::endl;
	for (int i = 0; i < size; ++i)
	{
		tree.deleteData(datas + i);
		if (!checkAVLTree(tree))
		{
			std::cout << "avl-tree delete exception" << std::endl;
			return 0;
		}
	}
	std::cout << "test avl tree successful!" << std::endl;
}

int main()
{
	int test_times = 100;
	for (int i = 0; i < test_times; ++i)
		unit_test();
	std::cout << "test over!" << std::endl;
	char p;
	std::cin >> p;
}
